//////////////////////////////////////////////////////////////////
// (c) copyright 2010- by jeongnim kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications 
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// supported by 
//   national center for supercomputing applications, uiuc
//   materials computation center, uiuc
//////////////////////////////////////////////////////////////////
// -*- c++ -*-
#ifndef BW_DNS_OPERATIONS_H
#define BW_DNS_OPERATIONS_H
#include <complex>
#if defined(ADD_)
#define do_rk4 do_rk4_
#define vxyz3 vxyz3_
#define vxy3z vxy3z_
#define v3xyz v3xyz_
#endif

//set it here
#ifndef INLINE0
#define INLINE0 inline
#endif

#ifndef INLINE1
#define INLINE1 inline
#endif

// declare fortran functions defined in dns_fops.f
#if defined(ENABLE_FORTRAN)
extern "C"
{
  /** equivalent RK::apply
   */
  void do_rk4(double* a,double* b,double* c, double *d
      , const double* d1, const double* d2, int* ntot);

  /** called by vorticity_3xyz_f
   */
  void v3xyz(std::complex<double>* target, const std::complex<double>* source
      ,const int* n1ind, const int* n1
      ,const int* n2ind, const int* n2
      ,const int* n3ind, const int *n3);

  /** called by vorticity_xy3z_f
   */
  void vxy3z(std::complex<double>* target, const std::complex<double>* source
      ,const int* n1ind,const  int* n1,const  int* n2ind, const int* n2, const int* n3ind, const int *n3);

  /** called by vorticity_xyz3_f
   */
  void vxyz3(std::complex<double>* target, const std::complex<double>* source
      ,const int* n1i, const int* n1f, const int* n2i, const int* n2f, const int* n3i, const int *n3f,const  int *ntot);
}
#endif


/** DNS operations are defined in dns_ops namespace
 */
namespace dns_ops
{
  /** implement 4-th order RK but ignore 4-d array
   * @param n equal to n_var*ny_2dz*nx_2dz*1,g_nz
   *
   * Original F9x routine looks like
   * \code
   do n=1,n_var
   do j=1,ny_2dz
   do i=1,nx_2dz
   do k=1,g_nz
   Q(k,i,j,n)=Q(k,i,j,n)+delt*rhs(k,i,j,n)/3
   Q_tmp(k,i,j,n)=Q_old(k,i,j,n) + delt*rhs(k,i,j,n)/2
   enddo
   enddo
   enddo
   enddo
   \endcode
   exactly the same operations as the original
   */
  template<typename T, typename RT>
  INLINE0 static void 
    apply(T* restrict Q,  T* restrict Q_tmp, T* restrict Q_old
        , const T* restrict rhs, RT one_third, RT delta, int n)
    {
      for(int i=0; i<n; ++i) 
      {
        Q[i]=one_third*rhs[i];
        Q_tmp[i]=Q_old[i]+delta*rhs[i];
      }
    }


  /** apply = apply1+apply2
   */
  template<typename T, typename RT>
  INLINE0 static void 
    apply1(T* restrict Q, const T* restrict rhs, RT one_third, int n)
    {
      for(int i=0; i<n; ++i) Q[i]+=one_third*rhs[i];
    }

  template<typename T, typename RT>
  INLINE0 static void 
    apply2(T* restrict Q, const T* restrict Q_old, const T* restrict rhs, RT delta, int n)
    {
      for(int i=0; i<n; ++i) Q[i] = Q_old[i] + delta*rhs[i];
    }

  /** aux function to evaluate vorticity_x in k-space for 3xyz storage
   */
  template<typename T, typename T2, typename ST>
  INLINE1 void
    vorticity_x(T* restrict x, const T* restrict y, const T* restrict z
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        for(int j=0; j<n2; ++j)
        {
          T jg=n2_map[j];
//#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            T kg=n3_map[k];
            x[ijk  ] = jg*z[ijk+1]-kg*y[ijk+1];
            x[ijk+1] = kg*y[ijk  ]-jg*z[ijk  ];
          }
        }
      }
    }

  /** aux function to evaluate vorticity_y in k-space for 3xyz storage
   */
  template<typename T, typename T2, typename ST>
  INLINE1 void
    vorticity_y(T* restrict y, const T* restrict z, const T* restrict x
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        T ig=n1_map[i];
        for(int j=0; j<n2; ++j)
        {
//#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            T kg=n3_map[k];
            y[ijk  ] = kg*x[ijk+1]-ig*z[ijk+1];
            y[ijk+1] = ig*z[ijk  ]-kg*x[ijk  ];
          }
        }
      }
    }

  /** aux function to evaluate vorticity_z in k-space for 3xyz storage
   */
  template<typename T, typename T2, typename ST>
  INLINE1 void
    vorticity_z(T* restrict z, const T* restrict x, const T* restrict y
        , const T2* restrict n1_map, ST n1
        , const T2* restrict n2_map, ST n2
        , const T2* restrict n3_map, ST n3
        )
    {
      for(int i=0,ijk=0; i<n1; ++i)
      {
        T ig=n1_map[i];
        for(int j=0; j<n2; ++j)
        {
          T jg=n2_map[j];
//#pragma unroll(16)
          for(int k=0; k<n3; ++k, ijk+=2)
          {
            z[ijk  ] = ig*y[ijk+1]-jg*x[ijk+1];
            z[ijk+1] = jg*x[ijk  ]-ig*y[ijk  ];
          }
        }
      }
    }

  /** evaluate the vorticity \f$ {\bf k} \times {\bf v}\f$ for 3xyz storage
   */
  template<typename T, typename T2>
  INLINE0 void
  vorticity_3xyz(T* restrict target, const T* restrict source
        , const T2* restrict n1_map, int n1
        , const T2* restrict n2_map, int n2
        , const T2* restrict n3_map, int n3
        )
    {
      int offset=n1*n2*n3*2;
      //this can be done independently
      vorticity_x(target         ,source+offset  ,source+2*offset,n1_map,n1,n2_map,n2,n3_map,n3);
      vorticity_y(target+offset  ,source+2*offset,source         ,n1_map,n1,n2_map,n2,n3_map,n3);
      vorticity_z(target+offset*2,source         ,source+offset  ,n1_map,n1,n2_map,n2,n3_map,n3);
    }


  template<typename T, typename T2>
  INLINE1 void vor_xyz(T* restrict x, const T* restrict y, const T* restrict z, T  jg
      , const T2* kmap, int n3)
  {
    for(int k=0,kk=0; k<n3; ++k, kk+=2)
    {
      T kg=kmap[k];
      x[kk  ] = jg*z[kk+1]-kg*y[kk+1];
      x[kk+1] = kg*y[kk  ]-jg*y[kk  ];
    }
  }

  template<typename T, typename T2>
  INLINE1 void vor_yzx(T* restrict y, const T* restrict z, const T* restrict x, T  ig
      , const T2* kmap, int n3)
  {
    for(int k=0,kk=0; k<n3; ++k, kk+=2)
    {
      T kg=kmap[k];
      y[kk  ] = kg*x[kk+1]-ig*z[kk+1];
      y[kk+1] = ig*z[kk  ]-kg*x[kk  ];
    }
  }

  template<typename T>
  INLINE1 void vor_zxy(T* restrict z, const T* restrict x, const T* restrict y, T  ig, T jg
      , int n3)
  {
    for(int k=0,kk=0; k<n3;++k,kk+=2)
    {
      z[kk  ] = ig*y[kk+1]-jg*x[kk+1];
      z[kk+1] = jg*x[kk  ]-ig*y[kk  ];
    }
  }

  /** target and source (n1,n2,3,n3,2)
   */
  template<typename T, typename T2>
  INLINE0 void
  vorticity_xy3z(T* restrict target, const T* restrict source
        , const T2* restrict n1_map, int n1
        , const T2* restrict n2_map, int n2
        , const T2* restrict n3_map, int n3
        )
    {
      for(int i=0; i<n1; ++i)
      {
        T ig=n1_map[i];
        for(int j=0; j<n2; ++j)
        {
          T jg=n2_map[j];
          int offset_x=(i*n2+j)*n3*6;
          int offset_y=offset_x+n3*2;
          int offset_z=offset_y+n3*2;
          vor_xyz(target+offset_x, source+offset_y, source+offset_z, jg, n3_map, n3);
          vor_yzx(target+offset_y, source+offset_z, source+offset_x, ig, n3_map, n3);
          vor_zxy(target+offset_z, source+offset_x, source+offset_y, ig, jg, n3);
        }
      }
    }


  template<typename T> inline void 
    vorticity_xy3z_f(std::complex<T>* restrict target, const std::complex<T>* restrict source
        , const int* restrict n1_map, int n1
        , const int* restrict n2_map, int n2
        , const int* restrict n3_map, int n3
        )
    {
      int nxyz3=n1*n2*n3*3;
      vxy3z(target,source,n1_map,&n1,n2_map,&n2,n3_map,&n3);
    }

  template<typename T> inline void 
    vorticity_xyz3_f(std::complex<T>* restrict target, const std::complex<T>* restrict source
        , const int* restrict n1_map, int n1
        , const int* restrict n2_map, int n2
        , const int* restrict n3_map, int n3
        )
    {
      int nxyz3=n1*n2*n3*3;
      vxyz3(target,source,n1_map,&n1,n2_map,&n2,n3_map,&n3,&nxyz3);
    }

  template<typename T> inline void 
    vorticity_3xyz_f(std::complex<T>* restrict target, const std::complex<T>* restrict source
        ,const int* restrict n1_map, int n1
        ,const int* restrict n2_map, int n2
        ,const int* restrict n3_map, int n3
        )
    {
      v3xyz(target,source,n1_map,&n1,n2_map,&n2,n3_map,&n3);
    }

  /**  target = outerproduct(target,rhs)
   * @param target  velocity in real space
   * @param rhs  vorticity in real spac
   * @param n total dimension
   * @param vorave average vorticity_z
   * @param ensave energy
   * @param helave \f$\sum dot(target,rhs)\f$
   */
  template<typename T>
    inline void
    force_xyz3(T* restrict target, const T* restrict rhs, int n
        )
        //, T& vorave, T& ensave, T& helave)
    {
      for(int i=0,ijk=0; i<n; ++i,ijk+=3)
      {
        T vorx=*rhs++; T vory=*rhs++; T vorz=*rhs++;
        //vorave+=vorz;
        //ensave+=vorx*vorx+vory*vory+vorz*vorz;
        //helave+=qx*vorx+qy*vory+qz*vorz;
        T qx=target[ijk]; T qy=target[ijk+1]; T qz=target[ijk+2];
        target[ijk  ] =qy*vorz-qz*vory;
        target[ijk+1] =qz*vorx-qx*vorz;
        target[ijk+2] =qx*vory-qy*vorz;
      }
    }

  template<typename T>
    inline void
    force_3xyz(T* restrict target, const T* restrict rhs, int n
        )
        //, T& vorave, T& ensave, T& helave)
    {
      int ix=0, iy=n, iz=2*n;
      for(int i=0; i<n; ++i,++ix,++iy,++iz)
      {
        T vorx=rhs[ix]; T vory=rhs[iy]; T vorz=rhs[iz];
        T qx=target[ix]; T qy=target[iy]; T qz=target[iz];
        //vorave+=vorz;
        //ensave+=vorx*vorx+vory*vory+vorz*vorz;
        //helave+=qx*vorx+qy*vory+qz*vorz;
        target[ix]=qy*vorz-qz*vory;
        target[iy]=qz*vorx-qx*vorz;
        target[iz]=qx*vory-qy*vorz;
      }
    }

  template<typename T>
    inline void outerproduct(const T* restrict rhsx, const T* restrict rhsy, const T* restrict rhsz
        ,T* restrict tx, T* restrict ty, T* restrict tz , size_t n)
    {
      for(size_t i=0; i<n;++i)
      {
        T vorx=rhsx[i]; 
        T vory=rhsy[i]; 
        T vorz=rhsz[i];
        T qx=tx[i]; 
        T qy=ty[i]; 
        T qz=tz[i];
        tx[i]=qy*vorz-qz*vory;
        ty[i]=qz*vorx-qx*vorz;
        tz[i]=qx*vory-qy*vorz;
      }
    }

  /** out += mkk * in
   * @param in input array 
   * @param out output array
   * @param mkk pre-computed laplacian factor
   * @param n size of mkk
   *
   * Note that in/out are complex data, i.e., in[2*n], out[2*n]
   */
  template<typename T>
    inline void multadd(const T* restrict in, const T* restrict mkk
        , T* restrict out, size_t n)
    {
      for(size_t i=0,i2=0; i<n;++i,i2+=2) 
      {
        out[i2] += mkk[i]*in[i2];
        out[i2+1] += mkk[i]*in[i2+1];
      }
    }

  /** target and source (n1_max,n2_max,n3_max,3)
   *
   * @param target new values of complex type
   * @param source old values of complex type
   * @param n1_first value of the first index
   * @param n1_last size of the first index
   * @param n2_first value of the second index
   * @param n2_last size of the second index
   * @param n3_first value of the third index
   * @param n3_last size of the third index
   */
  template<typename T>
  inline void
  vorticity_xyz3(std::complex<T>* restrict target, const std::complex<T>* restrict source
        , const int* restrict n1_map, int n1
        , const int* restrict n2_map, int n2
        , const int* restrict n3_map, int n3
        )
    {
      const T pi2=1.1;//2*PI;
      for(int i=0,ijk=0; i<n1; ++i)
      {
        std::complex<T> ig(0,pi2*n1_map[i]);
        for(int j=0; j<n2; ++j)
        {
          std::complex<T> jg(0,pi2*n2_map[j]);
          for(int k=0; k<n3; ++k,ijk+=3)
          {
            std::complex<T> kg(0,pi2*n3_map[k]);
            //target[ijk]=outprodct(k,source[ijk]);
            target[ijk  ]= jg*source[ijk+2]-kg*source[ijk+1];
            target[ijk+1]= kg*source[ijk  ]-ig*source[ijk+2];
            target[ijk+2]= ig*source[ijk+1]-jg*source[ijk  ];
          }
        }
      }
    }

  /** target and source (n0.size(),n1.size(),n2.size(),3,2)
   */
  template<typename T>
  inline void
  vorticity_xyz3(T* restrict target, const T* restrict source
        , const int* restrict n1_map, int n1
        , const int* restrict n2_map, int n2
        , const int* restrict n3_map, int n3
        )
    {
      const T pi2=1.1;//2*PI;
      enum {indx_r=0, indx_i, indy_r, indy_i, indz_r, indz_i};
      for(int i=0,ijk=0; i<n1; ++i)
      {
        T ig=pi2*n1_map[i];
        for(int j=0; j<n2; ++j)
        {
          T jg=pi2*n2_map[j];
          for(int k=0; k<n3; ++k, ijk+=6)
          {
            T kg=pi2*n3_map[k];
            target[ijk] = jg*source[ijk+indz_i]-kg*source[ijk+indy_i]; 
            target[ijk+1] = kg*source[ijk+indy_r]-jg*source[ijk+indz_r];
            target[ijk+2] = kg*source[ijk+indx_i]-ig*source[ijk+indz_i]; 
            target[ijk+3] = ig*source[ijk+indz_r]-kg*source[ijk+indx_r];
            target[ijk+4] = ig*source[ijk+indy_i]-jg*source[ijk+indx_i]; 
            target[ijk+5] = jg*source[ijk+indx_r]-ig*source[ijk+indy_r];
          }
        }
      }
    }

}

#if !defined(ENABLE_FORTRAN)
  template<typename T, typename RT>
    inline void do_rk4(T* a,T* b,T* c, T *d , RT* d1, RT* d2, int* ntot)
    {
      dns_ops::apply(a,b,c,d,*d1,*d2,*ntot);
    }
#endif

#endif
/***************************************************************************
 * $RCSfile: codelets.hpp,v $   $Author: jnkim $
 * $Revision: 1.2 $   $Date: 2005/12/11 15:28:52 $
 * $Id: codelets.hpp,v 1.2 2005/12/11 15:28:52 jnkim Exp $
 ************************************************************************/
